import os
import sys
from collections import defaultdict

from openpyxl import load_workbook

# 0-based index
CFO_ID_COL_IDX = 2
COMP_ID_COL_IDX = 5
YEAR_COL_IDX = 6
GENDER_COL_IDX = 7
# 1-based index
Y_COL_IDX = 13
MARK_COL_IDX = {
    'Male_to_Male': 14,
    'Male_to_Female': 15,
    'Female_to_Male': 16,
    'Female_to_Female': 17,
}
CHANGE_NOTE = 'CFO Change'


def find_split(l, only_first_change:bool, target_change_type:str):
    change_count = 0
    for i in range(len(l)):
        before_cfo = l[i - 1]['cfo'] if i > 0 else None
        after_cfo = l[i + 1]['cfo'] if i + 1 < len(l) else None
        if before_cfo == after_cfo:
            continue

        before_counting = 0
        before_bounder = i - 1
        for j in range(i - 1, -1, -1):  # [i-1,0]
            if before_cfo == l[j]['cfo'] and CHANGE_NOTE not in str(l[j].get(target_change_type,'')):
                before_counting += 1
                before_bounder = j
            else:
                break
        if before_counting < 2:
            continue

        after_counting = 0
        after_bounder = i + 1
        for k in range(i + 1, len(l), 1):  # [i+1, len(l)-1]
            if after_cfo == l[k]['cfo'] and CHANGE_NOTE not in str(l[k].get(target_change_type,'')):
                after_counting += 1
                after_bounder = k
            else:
                break

        if before_counting >= 2 and after_counting >= 2 and l[i]['cfo'] == after_cfo:  # AAABBBB is OK, but AAAABBB not
            changing_type = l[i - 1]['gender'].capitalize() + '_to_' + l[i + 1]['gender'].capitalize()
            before_notation = 0
            for a in range(i - 1, before_bounder - 1, -1):
                l[a][changing_type] = before_notation

            if changing_type == target_change_type:
                change_count += 1
            l[i][changing_type] = CHANGE_NOTE + str(change_count)

            after_notation = 1
            for a in range(i + 1, after_bounder + 1):
                l[a][changing_type] = after_notation
            if only_first_change:
                break
    return l


def main(file_name: str, only_first_change: bool, change_type:str):
    input_file = os.path.join('dataset', file_name)
    out_file = os.path.join('result',
                            '{}_{}{}'.format(file_name[:-5], 'first' if only_first_change else 'all', file_name[-5:]))

    wb = load_workbook(filename=input_file)
    ws = wb['Final']
    cnt = 0
    group_records = defaultdict(list)  # key is company, values are cfo records
    is_header = True
    for i, r in enumerate(ws.rows):
        if is_header:
            is_header = False
            continue
        cfo = r[CFO_ID_COL_IDX].value  # starting from 0
        comp = r[COMP_ID_COL_IDX].value
        year = r[YEAR_COL_IDX].value
        gender = r[GENDER_COL_IDX].value
        group_records[comp].append({
            'rn': i + 1,
            'cfo': cfo,
            'comp': comp,
            'year': year,
            'gender': gender
        })
        cnt += 1

    print("load done")
    for k, v in group_records.items():
        find_split(v, only_first_change, change_type)

    for cfo_list in group_records.values():
        for cfo in cfo_list:
            row_no = cfo['rn']
            if change_type in cfo.keys():
                mark_column_index = MARK_COL_IDX[change_type]
                ws.cell(row=row_no, column=mark_column_index).value = cfo[change_type]
                ws.cell(row=row_no, column=Y_COL_IDX).value = cfo['year']
    wb.save(out_file)


if __name__ == '__main__':
    if len(sys.argv) > 1 and sys.argv[1] == '-f':
        only_first_change = True
    else:   # default to find all changes
        only_first_change = False

    change_type = 'Male_to_Male'
    filename = 'CFO Change 2009-2019-V20211207_only_male_to_male.xlsx'
    main(filename, only_first_change, change_type)

    change_type = 'Female_to_Male'
    filename = 'CFO Change 2009-2019-V20211207_only_female_to_male.xlsx'
    main(filename, only_first_change, change_type)

    change_type = 'Male_to_Female'
    filename = 'CFO Change 2009-2019-V20211207_only_male_to_female.xlsx'
    main(filename, only_first_change, change_type)